from enum import Enum

from pydantic import BaseModel


class Metrics(BaseModel):
    train_loss: float
    train_accuracy: float
    val_loss: float
    val_accuracy: float


class Framework(Enum):
    PyTorch = "torch"
    Tensorflow = "tensorflow"
    XGBoost = "xgboost"
    LightGBM = "lightgbm"


class Artifact(BaseModel):
    name: str
    metrics: Metrics
    framework: Framework
    archive_link: str
    job_url: str
    grafana_url: str


class ExecutionFlow(BaseModel):
    flows: list[str]


class ArtifactInDB(BaseModel):
    id: str | None
    name: str
    uuid: str
    artifact_info: str
